import torch
import torch.nn as nn
import torch.nn.functional as F

CONV_KWARGS = dict(kernel_size=3, padding=1)


class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, padding=1):
        super().__init__()
        self.depthwise = nn.Conv2d(
            in_channels,
            in_channels,
            kernel_size=kernel_size,
            padding=padding,
            groups=in_channels,
        )
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x


class AttentionBlock(nn.Module):
    def __init__(self, F_g, F_l, F_int):
        super().__init__()
        self.W_g = nn.Sequential(
            nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int),
        )
        self.W_x = nn.Sequential(
            nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int),
        )
        self.psi = nn.Sequential(
            nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid(),
        )

    def forward(self, g, x):
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = self.psi(F.relu(g1 + x1))
        return x * psi


class EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        # self.conv1 = DepthwiseSeparableConv(in_channels, out_channels)
        self.conv1 = nn.Conv2d(in_channels, out_channels, **CONV_KWARGS)
        self.norm1 = nn.BatchNorm2d(out_channels)
        # self.conv2 = DepthwiseSeparableConv(out_channels, out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, **CONV_KWARGS)
        self.norm2 = nn.BatchNorm2d(out_channels)
        self.pool = nn.MaxPool2d(2)
        self.activation = nn.ReLU()

    def forward(self, x):
        x = self.activation(self.norm1(self.conv1(x)))
        x = self.activation(self.norm2(self.conv2(x)))
        p = self.pool(x)
        return x, p


class DecoderBlock(nn.Module):
    def __init__(self, in_channels, skip_channels, out_channels):
        super().__init__()
        # 更新了这里的in_channels，加上了skip_channels
        self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
        self.attention = AttentionBlock(out_channels, skip_channels, skip_channels // 2)
        # self.conv1 = DepthwiseSeparableConv(out_channels + skip_channels, out_channels)
        self.conv1 = nn.Conv2d(
            out_channels + skip_channels, out_channels, **CONV_KWARGS
        )
        self.norm1 = nn.BatchNorm2d(out_channels)
        # self.conv2 = DepthwiseSeparableConv(out_channels, out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, **CONV_KWARGS)
        self.norm2 = nn.BatchNorm2d(out_channels)
        self.activation = nn.ReLU()

    def forward(self, x, skip):
        x = self.up(x)
        skip = self.attention(x, skip)
        # 在合并前进行上采样
        x = torch.cat([x, skip], dim=1)
        x = self.activation(self.norm1(self.conv1(x)))
        x = self.activation(self.norm2(self.conv2(x)))
        return x


class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, min_channels=64):
        super().__init__()
        self.enc1 = EncoderBlock(in_channels, min_channels)
        self.enc2 = EncoderBlock(min_channels, min_channels * 2)
        self.enc3 = EncoderBlock(min_channels * 2, min_channels * 4)
        self.enc4 = EncoderBlock(min_channels * 4, min_channels * 8)

        self.center = DepthwiseSeparableConv(min_channels * 8, min_channels * 16)

        self.dec4 = DecoderBlock(min_channels * 16, min_channels * 8, min_channels * 8)
        self.dec3 = DecoderBlock(min_channels * 8, min_channels * 4, min_channels * 4)
        self.dec2 = DecoderBlock(min_channels * 4, min_channels * 2, min_channels * 2)
        self.dec1 = DecoderBlock(min_channels * 2, min_channels, min_channels)

        self.final_conv = nn.Conv2d(min_channels, out_channels, kernel_size=1)

    def forward(self, x):
        x1, p1 = self.enc1(x)
        x2, p2 = self.enc2(p1)
        x3, p3 = self.enc3(p2)
        x4, p4 = self.enc4(p3)

        center = self.center(p4)

        d4 = self.dec4(center, x4)
        d3 = self.dec3(d4, x3)
        d2 = self.dec2(d3, x2)
        d1 = self.dec1(d2, x1)

        out = self.final_conv(d1)
        return torch.sigmoid(out)
